#include "TrafficiDentify.h"

//红绿灯的hsv阈值，轻易不要动
int hsv_red[2][3] = {
    {160, 43, 178},
    {180, 255, 255},
};
int hsv_yellow[2][3] = {
    {11, 100, 178},
    {40, 255, 255},
};
int hsv_green[2][3] = {
    {40, 50, 178},
    {97, 255, 255},
};

bool areasortfun(RotatedRect first, RotatedRect second){
    return first.size.area() > second.size.area();
}

bool y_sortfun_inv(RotatedRect first, RotatedRect second){
    return first.center.y < second.center.y;
}
Rect assert_boundRect(Rect in, Mat roi){
    Rect temp;
    if (in.x < 0){
        temp.x = 0;
        in.x = 0;
    }
    else{
        temp.x = in.x;
    }

    if (in.y < 0){
        temp.y = 0;
        in.y = 0;
    }
    else{
        temp.y = in.y;
    }

    temp.x = in.x > roi.cols ? roi.cols : in.x;
    temp.y = in.y > roi.rows ? roi.rows : in.y;

    temp.width = in.width < 0 ? 0 : in.width;
    temp.height = in.height < 0 ? 0 : in.height;
    if (in.x + in.width > roi.cols){
        temp.width = roi.cols - temp.x;
    }
    if (in.height + in.y > roi.rows){
        temp.height = roi.rows - temp.y;
    }
    return temp;
}

vector<RotatedRect> find_area(Mat& img_gray, RotatedRect_ptr sortfun, int areathreshold = 300, int maxarea = 16000){
    std::vector<std::vector<Point>> contours;
    cv::findContours(img_gray, contours, RETR_EXTERNAL, CHAIN_APPROX_NONE);

    vector<RotatedRect> minRect;
    for (int i = 0; i < (int)contours.size(); i++){
        RotatedRect tempRect;
        tempRect = minAreaRect(Mat(contours[i]));
        if (tempRect.size.area() > areathreshold  && tempRect.size.area() < maxarea){
            minRect.push_back(tempRect);
        }
    }
    sort(minRect.begin(), minRect.end(), sortfun);
    return minRect;
}
vector<RotatedRect> find_Lightarea(Mat& img_gray, RotatedRect_ptr sortfun,int thr = 5, int areathreshold = 150){
    std::vector<std::vector<Point>> contours;
    cv::findContours(img_gray, contours, RETR_EXTERNAL, CHAIN_APPROX_NONE);

    vector<RotatedRect> minRect;
    for (int i = 0; i < (int)contours.size(); i++){
        RotatedRect tempRect;
        tempRect = minAreaRect(Mat(contours[i]));
        if (tempRect.size.area() > areathreshold){
            minRect.push_back(tempRect);
        }
    }
    sort(minRect.begin(), minRect.end(), sortfun);
    return minRect;
}

void draw_RotatedRect(Mat img, RotatedRect rect, Scalar color){
    Point2f rect_points[4];
    rect.points(rect_points);
    for (int i = 0; i < 4; i++){
        rect_points[i].x = rect_points[i].x < 0 ? 0 : rect_points[i].x;
        rect_points[i].y = rect_points[i].y < 0 ? 0 : rect_points[i].y;
    }
    for (int i = 0; i < 2; i++){
        for (int j = 0; j < 4; j++){
            if (img.channels() == 3)
                line(img, rect_points[j], rect_points[(j + 1) % 4], color, 2, 8);
            else{
                line(img, rect_points[j], rect_points[(j + 1) % 4], color, 2, 8);
            }
        }
    }
}

int findZebarCrossing(Mat inputMat, vector<RotatedRect>& roi){
    Mat srcImg;
    float radio = 0.25f;
    //取消掉图像的上四分之一处，如果上面那部分有斑马线，那也只能说明斑马线离摄像头很远，不是眼前的斑马线
    inputMat(Rect(0,
        cvCeil(inputMat.rows * radio),
        inputMat.cols,
        inputMat.rows - cvCeil(inputMat.rows * radio)))
        .copyTo(srcImg);

    Mat imgBin = Mat::zeros(srcImg.rows, srcImg.cols, CV_8UC1);
    Mat grayImg;
    cvtColor(srcImg, grayImg, COLOR_BGR2GRAY);

    //gapSize的意识是把图像横切成几份，对每一部分进行自适应二值化，因为斑马线是白色的会比较亮，
    int gapSize = 6;
    int gap = grayImg.rows / gapSize;
    for (int i = 0; i < gapSize - 1; i++){
        threshold(grayImg(Rect(0, i * gap, grayImg.cols, gap)), imgBin(Rect(0, i * gap, grayImg.cols, gap)), 0, 255, THRESH_OTSU);
    }
    threshold(grayImg(Rect(0, (gapSize - 1) * gap, grayImg.cols, grayImg.rows - (gapSize - 1) * gap)), 
        imgBin(Rect(0, (gapSize - 1) * gap, grayImg.cols, grayImg.rows - (gapSize - 1) * gap)), 0, 255, THRESH_OTSU);
    //中值滤波，用来滤除椒盐噪声，因为路面是柏油面，会有很多反射太阳光的小光点，用这个把这些小光点滤除掉
    medianBlur(imgBin, imgBin, 3);
    //形态学运算中的开运算，用来滤除图像中的噪声
    Mat element = getStructuringElement(MORPH_RECT, Size(4, 4)); //得到一个4×4的运算核
    morphologyEx(imgBin, imgBin, MORPH_OPEN, element);
    //imshow("Bin", imgBin);
    //寻找二值化图像中的连通域，并且使用最小的可以覆盖该连通域的矩形框给包围起来然后以面积从大到小排序
    vector<RotatedRect> minRect = find_area(imgBin, areasortfun);
    //cout << "size:" << minRect.size() << endl;
    //只要前二十个连通域，也是用来滤除过小的连通域
    if (minRect.size() > 20){
        while (minRect.size() > 20){
            minRect.pop_back();
        }
    }
    if (minRect.size() == 0){
        return 0;
    }

    //两个数组，用来存储每个矩形框的:
    //w_h_radio是宽高比，dis_center是这个矩形框的中心点和图像中心点的差值并且将改值归一化到1-255
    //
    /*
        以w_h_radio为例，假如里面存储的是如下数据
        1,2,3,3,3,4,2,3,4,3,3，2,2,3,4,2,3，2，3
        因为斑马线上的横条是类似的，并且都成堆出现，但是噪声和其他的干扰因素总是单个出现的
        就比如上面2和3出现的次数非常多
        那么此时我们可以认为宽高比为2和3的矩形框包围的是斑马线，1和4就是干扰因素
        同样dis_center也是一样的，斑马线与图像中线的偏差也是总是成堆出现的

    */

    Mat w_h_radio = Mat::zeros(1, minRect.size(), CV_8UC1);
    Mat dis_center = Mat::zeros(1, minRect.size(), CV_8UC1);

    for (int i = 0; i < (int)minRect.size(); i++){
        w_h_radio.at<uchar>(0, i) = minRect[i].boundingRect().width / minRect[i].boundingRect().height;
        dis_center.at<uchar>(0, i) = cvCeil((abs(imgBin.cols * 0.5 - minRect[i].center.x) / (imgBin.cols * 0.5)) * 255);
    }

    //对上面的计算出来的参数使用 最大类间方差的方法是的使得出现次数比较多的数据凸显出来，比如上面例子中的2和3
    threshold(w_h_radio, w_h_radio, 0, 255, THRESH_OTSU);
    threshold(dis_center, dis_center, 0, 255, THRESH_OTSU);
    for (int i = 0; i < (int)minRect.size(); i++){
        if (w_h_radio.at<uchar>(0, i) && (!dis_center.at<uchar>(0, i)) && (minRect[i].boundingRect().width > srcImg.cols * 0.2)){
            //把选择出来的斑马线假如roi里面
            roi.push_back(minRect[i]);
        }
    }

    if (roi.size() < 6){
        return 0; //代表当前没有斑马线
    }

    //因为我们一直是在图像的下四分之三处搜索，现在将斑马线的区域映射到原来那个完整的图像中
    for (int i = 0; i < (int)roi.size(); i++){
        roi[i].center.y += inputMat.rows - srcImg.rows;
    }
    return 1;
}

int judgeLights(Mat roi, int color, RotatedRect color_area){
    if (roi.empty())
        return -1;
    int x = color_area.center.x - color_area.boundingRect().x;
    Mat gray;
    Mat binary_dark;

    cvtColor(roi, gray, COLOR_BGR2GRAY);
    //红绿灯的上或者下总是有没亮的灯，提取这部分黑色区域来验证
    threshold(gray, binary_dark, 50, 255, THRESH_BINARY_INV);
    //Mat element = getStructuringElement(MORPH_RECT, Size(2, 2)); //得到一个2×2的运算核
    //中值滤波
    medianBlur(binary_dark, binary_dark, 3);
    //imshow("binary_dark", binary_dark);

    //以高度排序，红绿灯大概率是所选择区域里面最高的拿一个，虽然这样会降低红绿灯的检测率，但是极大程度上抑制了干扰
    vector<RotatedRect> areas_dark, areas_light;
    areas_dark = find_Lightarea(binary_dark, y_sortfun_inv, 5, 100);
    if (color == red || color == green){ //红灯，绿灯和黄灯处理方式不一样，分开处理
        for (int i = 0; i < (int)areas_dark.size(); i++){
            int radio = areas_dark[i].size.height / areas_dark[i].size.width; //高宽比，红绿灯的高除以宽大约在2-3之间
            if (
                i == 0 //只要第一个，就是最高的那一个
                && abs(areas_dark[i].center.x - x) < areas_dark[i].size.width && radio >= 2 && radio <= 4
                ){
                return color;
            }
        }
    }
    else if (color == yellow){
        for (int i = 0; i < (int)areas_dark.size(); i++){
            int radio = areas_dark[i].size.height / areas_dark[i].size.width;
            if (
                i == 0 
                && abs(areas_dark[i].center.x - x) < areas_dark[i].size.width && radio >= 1 && radio <= 2
                ){
                return color;
            }
        }
    }
    return -1;
}

int findSignal(Mat inputMat){
    Mat imgHSV;
    //红绿灯只在图像上面那部分进行搜索，如果红绿灯出现在图像下面那部分，那就代表这个红绿灯不是我们要看的马路对面的那个红绿灯
    Mat srcImg = inputMat(Rect(0, 0, inputMat.cols, cvCeil(inputMat.rows / 2.5))).clone(); //cvCeil()：返回不小于参数的最小整数值，即向上取整；
    Mat redImg = Mat::zeros(Size(srcImg.cols, srcImg.rows), CV_8UC1);
    Mat greenImg = Mat::zeros(Size(srcImg.cols, srcImg.rows), CV_8UC1);
    Mat yellowImg = Mat::zeros(Size(srcImg.cols, srcImg.rows), CV_8UC1);
    cvtColor(srcImg, imgHSV, COLOR_BGR2HSV); //从RGB空间转换成HSV空间，方便颜色的识别

    //分别对红，黄(实际是橙色)，绿三色进行提取
    inRange(imgHSV, Scalar(hsv_red[low][HSV_H], hsv_red[low][HSV_S], hsv_red[low][HSV_V]), 
        Scalar(hsv_red[high][HSV_H], hsv_red[high][HSV_S], hsv_red[high][HSV_V]), redImg);

    inRange(imgHSV, Scalar(hsv_green[low][HSV_H], hsv_green[low][HSV_S], hsv_green[low][HSV_V]), 
        Scalar(hsv_green[high][HSV_H], hsv_green[high][HSV_S], hsv_green[high][HSV_V]), greenImg);

    inRange(imgHSV, Scalar(hsv_yellow[low][HSV_H], hsv_yellow[low][HSV_S], hsv_yellow[low][HSV_V]), 
        Scalar(hsv_yellow[high][HSV_H], hsv_yellow[high][HSV_S], hsv_yellow[high][HSV_V]), yellowImg);

    //由于红绿灯一般在马路对面，光斑比较小，所以我们要对每个图像做一个膨胀
    Mat element = getStructuringElement(MORPH_RECT, Size(4, 4));
    dilate(greenImg, greenImg, element);
    dilate(yellowImg, yellowImg, element);
    dilate(redImg, redImg, element);

    /*Mat temp;
    srcImg.copyTo(temp, greenImg);
    imshow("green", temp);
    Mat temp_1;
    srcImg.copyTo(temp_1, redImg);
    imshow("red", temp_1);
    Mat temp_2;
    srcImg.copyTo(temp_2, yellowImg);
    imshow("yellow", temp_2);*/

    vector<RotatedRect> area_green = find_Lightarea(greenImg, areasortfun, 5,65);
    vector<RotatedRect> area_red = find_Lightarea(redImg, areasortfun, 5, 65);
    vector<RotatedRect> area_yellow = find_Lightarea(yellowImg, areasortfun, 5, 65);
    if (area_green.size() > 5 || area_red.size() > 5 || area_yellow.size() > 5){
        Mat ylocSelect = Mat::zeros(1, area_green.size() + area_red.size() + area_yellow.size(), CV_16UC1);
        int index = 0;
        for (int i = 0; i < (int)area_green.size(); i++, index++){
            ylocSelect.at<ushort>(0, index) = area_green[i].boundingRect().y;
        }
        for (int i = 0; i < (int)area_red.size(); i++, index++){
            ylocSelect.at<ushort>(0, index) = area_red[i].boundingRect().y;
        }
        for (int i = 0; i < (int)area_yellow.size(); i++, index++){
            ylocSelect.at<ushort>(0, index) = area_yellow[i].boundingRect().y;
        }
        normalize(ylocSelect, ylocSelect, 0, 255, NORM_MINMAX);//数据归一化  线性归一化
        Mat cats_8UC1 = Mat::zeros(ylocSelect.size(), CV_8UC1);
        for (int i = 0; i < ylocSelect.cols; i++){
            cats_8UC1.at<uchar>(0, i) = ylocSelect.at<ushort>(0, i);
        }
        threshold(cats_8UC1, cats_8UC1, 0, 255, THRESH_TRIANGLE);
        index = 0;
        for (auto it = area_green.begin(); it != area_green.end(); index++){
            if (cats_8UC1.at<uchar>(0, index)){
                it = area_green.erase(it);
            }
            else{
                it++;
            }
        }
        for (auto it = area_red.begin(); it != area_red.end(); index++){
            if (cats_8UC1.at<uchar>(0, index)){
                it = area_red.erase(it);
            }
            else{
                it++;
            }
        }
        for (auto it = area_yellow.begin(); it != area_yellow.end(); index++){
            if (cats_8UC1.at<uchar>(0, index)){
                it = area_yellow.erase(it);
            }
            else{
                it++;
            }
        }
    }

    for (int i = 0; i < (int)area_green.size(); i++){ //绿灯在下，通过找到的光斑，拓展到整个红绿灯
        area_green[i].center.y -= area_green[i].size.height;
        area_green[i].size.height *= 3;
        area_green[i].size.width *= 3;
    }
    for (int i = 0; i < (int)area_yellow.size(); i++){ //黄灯在中，通过找到的光斑，拓展到整个红绿灯
        area_yellow[i].size.height *= 3;
        area_yellow[i].size.width *= 3;
    }
    for (int i = 0; i < (int)area_red.size(); i++){ //红灯在上，通过找到的光斑，拓展到整个红绿灯
        area_red[i].center.y += area_red[i].size.height;
        area_red[i].size.height *= 3;
        area_red[i].size.width *= 3;
    }
    //Mat temp_3;
    int maxsize = 1000;
    for (int i = 0; i < (int)area_green.size(); i++){
        if (area_green[i].size.area() < maxsize && green == judgeLights(srcImg(assert_boundRect(area_green[i].boundingRect(), srcImg)).clone(), 
            green, area_green[i])){
            //temp_3 = srcImg(assert_boundRect(area_green[i].boundingRect(), srcImg)).clone();
            //imshow("temp_3", temp_3);
            return green;
        }
    }
    for (int i = 0; i < (int)area_yellow.size(); i++){
        if (area_yellow[i].size.area() < maxsize && yellow == judgeLights(srcImg(assert_boundRect(area_yellow[i].boundingRect(), srcImg)).clone(), 
            yellow, area_yellow[i])){
            //temp_3 = srcImg(assert_boundRect(area_yellow[i].boundingRect(), srcImg)).clone();
            //imshow("temp_3", temp_3);
            return yellow;
        }
    }
    for (int i = 0; i < (int)area_red.size(); i++){
        if (area_red[i].size.area() < maxsize && red == judgeLights(srcImg(assert_boundRect(area_red[i].boundingRect(), srcImg)).clone(),
            red, area_red[i])){
            //temp_3 = srcImg(assert_boundRect(area_red[i].boundingRect(), srcImg)).clone();
            //imshow("temp_3", temp_3);
            return red;
        }
    }
    return -1;
}